import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor, Compose
from PIL import Image


class MnistDataset(Dataset):
    def __init__(self, img_dir_path, label_txt_path, transform):
        self.img_dir_path = img_dir_path
        self.label_txt_path = label_txt_path
        self.transform = transform
        with open(label_txt_path, "r") as f:
            lines = f.readlines()
        self.img_paths = []
        self.labels = []
        for line in lines:
            data = line.split(" ")
            self.img_paths.append(data[0])
            self.labels.append(data[1])
        #with open(label_txt_path, "r") as f:
            # 使用列表推导式简化初始化过程
        #    self.img_paths, self.labels = zip(*[(line.split()[0], int(line.split()[1])) for line in f])

    def __getitem__(self, idx):
        img = Image.open(self.img_paths[idx])
        img = self.transform(img)
        label = int(self.labels[idx])
        return img, label

    def __len__(self):
        return len(self.labels)


if __name__ == "__main__":
    train_img_dir = "./data/MNIST/train"
    train_label_txt = "./data/MNIST/train.txt"
    test_img_dir = "./data/MNIST/test"
    test_label_txt = "./data/MNIST/test.txt"
    transform = Compose([ToTensor()])

    train_dataset = MnistDataset(img_dir_path=train_img_dir, label_txt_path=train_label_txt, transform=transform)
    test_dataset = MnistDataset(img_dir_path=test_img_dir, label_txt_path=test_label_txt, transform=transform)

    # print(len(train_dataset))
    # data=train_dataset[0]
    # img,label=data[0],data[1]
    # img = img * 255
    # img=img.numpy().astype(np.uint8).squeeze()
    # img=Image.fromarray(img)
    # print(label)
    # img.show()
    #
    # print(len(test_dataset))
    # data = test_dataset[0]
    # img, label = data[0], data[1]
    # data = train_dataset[0]
    # img, label = data[0], data[1]
    # img = img * 255
    # img = img.numpy().astype(np.uint8).squeeze()
    # img = Image.fromarray(img)
    # print(label)
    # img.show()

    train_dl = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_dl = DataLoader(test_dataset, batch_size=1024, shuffle=False)

    for i ,(imgs, labels)  in enumerate(train_dl):
        # print(imgs.shape)
        # print(imgs[0].shape)
        # print(list(imgs))
        # print(type(imgs))
        print(imgs.shape,labels.shape)
        break